from langchain.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain.chains.sequential import SequentialChain
from langchain.chains.llm import LLMChain

api_key = "sk-6S0PtpNia71gjcfwSsDPsJ9mGqsVPr2XRQzAx1dHbJS7RW4t"
api_base="https://chatapi.littlewheat.com/v1"

llm = ChatOpenAI(model="gpt-3.5-turbo",api_key=api_key ,base_url=api_base)

chat1 = ChatPromptTemplate.from_messages([
("system", "你是一位精通各领域知识的知名教授"),
        ("human", "请你先尽可能详细的解释一下：{knowledge}，并且{action}")
])

chain1 = LLMChain(llm=llm,prompt=chat1,output_key="chat1_output_key")

chat2 = ChatPromptTemplate.from_messages([
("system", "你非常善于提取文本中的重要信息，并做出简短的总结"),
        ("human", "这是针对一个提问完整的解释说明内容：{chat1_output_key}"),
        ("human", "请你根据上述说明，尽可能简短的输出重要的结论，请控制在100个字以内")
])
chain2=LLMChain(llm=llm,prompt=chat2,output_key="chat2_output_key")

chain3 = SequentialChain(chains=[chain1,chain2],
                         input_variables=["knowledge", "action"],
                         output_variables=["chat1_output_key", "chat2_output_key"],
                         verbose=True)
response=chain3.invoke({
"knowledge":"中国乒乓球水平怎么样？",
                       "action":"举一个实际的例子"
})
print(response)
print(response["chat1_output_key"])